import numpy as np
from LTE import constant as c
from LTE.phy import Phy

class PRInit:
    SEQUENCE_SEED_LEN = 31
    SEQUENCE_NC = 1600
    @staticmethod
    def step(state,g):
        f = g(state)
        f = (f & 1) << (PRInit.SEQUENCE_SEED_LEN - 1)
        state = (state >> 1) ^ f
        return state
    
    @staticmethod
    def step_x1(state):
        return PRInit.step(state,lambda x:x^(x>>3))
        
    @staticmethod
    def step_x2(state):
        return PRInit.step(state,lambda x:x^(x>>1)^(x>>2)^(x>>3))
    @staticmethod
    def init(f,i):
        for _ in range(PRInit.SEQUENCE_NC):
            i = f(i)
        return i

class PR:
    X1_Init = PRInit.init(PRInit.step_x1,1)
    X2_Init = [PRInit.init(PRInit.step_x2,1<<i) for i in range(PRInit.SEQUENCE_SEED_LEN)]
    @staticmethod
    def x2_init(seed):
        x2 = 0
        for i in range(PRInit.SEQUENCE_SEED_LEN):
            if ((seed >> i) & 1)==1:
                x2 ^= PR.X2_Init[i]
        return x2
    
    @staticmethod    
    def seq(seed,len):
        x1 = PR.X1_Init
        x2 = PR.x2_init(seed)
        pr = np.zeros(len,np.uint8)
        for n in range(len):
            pr[n] = (x1 ^ x2) & 1
            x1 = PRInit.step_x1(x1)
            x2 = PRInit.step_x2(x2)
        return pr

class Cell:
    def __init__(self,id,nof_prb=6,cp="normal"):
        self.id = id
        self.nof_prb = nof_prb
        self.cp = cp
        self.nof_ports = 4
        self.init()

    def init(self):
        self.pss_signal  = PSS(self.id%3)
        self.sss_signal0 = SSS(self.id//3,self.id%3,0)
        self.sss_signal5 = SSS(self.id//3,self.id%3,5)
        self.pilots = self.pilot()

    def x2_init(self,ns,l,p):
        N_cp = 1 if self.cp=="normal" else 0
        lp = c.refsignal_cs_nsymbol(l, self.cp, 2 * p)
        ci = 1024 * (7 * (ns + 1) + lp + 1) * (2 * self.id + 1) + 2 * self.id + N_cp
        return ci

    def pilot(self):
        cs_of_sym = c.refsignal_cs_nof_symbols(2 * self.nof_ports)
        pilots = np.zeros((2,c.LTE_NOF_SF_X_FRAME,4*cs_of_sym*self.nof_prb),np.complex64)
        for ns in range(c.LTE_NSLOTS_X_FRAME):
            for p in range(2):
                nsymbols = c.refsignal_cs_nof_symbols(2 * p) // 2
                for l in range(nsymbols):
                    c_init = self.x2_init(ns,l,p)
                    seq = np.array(PR.seq(c_init, 2 * 2 * c.LTE_MAX_PRB))
                    c_seq = (1.0-2.0*seq)*np.sqrt(0.5)
                    for i in range(2 * self.nof_prb):
                        idx = c.LTE_REFSIGNAL_PILOT_IDX(i, (ns % 2) * nsymbols + l, self.nof_prb)
                        mp  = i + c.LTE_MAX_PRB - self.nof_prb
                        pilots[p][ns//2][idx] = c_seq[2*mp] + 1j * c_seq[2*mp+1]
        return pilots
                    
    def cs_put(self,port_id):
        sf_symbols = np.zeros(
            ( c.LTE_NOF_SF_X_FRAME
            , 2*c.LTE_CP_NSYMB(self.cp)
            , self.nof_prb * c.LTE_NRE)
            , np.complex64)
        for sf in range(c.LTE_NOF_SF_X_FRAME):
            if port_id ==-1 and (sf == 0 or sf == 5):
                sss = self.sss_signal0.ofd if sf==0 else self.sss_signal5.ofd
                sf_symbols[sf,c.LTE_CP_NORM_NSYMB-1,5:5+62] = self.pss_signal.ofd[:]
                sf_symbols[sf,c.LTE_CP_NORM_NSYMB-2,5:5+62] = sss[:]
            for l in range(c.refsignal_cs_nof_symbols(port_id)):
                nsymbol = c.refsignal_cs_nsymbol(l, self.cp, port_id)
                fidx = (c.refsignal_cs_v(port_id, l) + (self.id % 6)) % 6
                for i in range(2 * self.nof_prb):
                    pidx = c.LTE_REFSIGNAL_PILOT_IDX(i, l, self.nof_prb)
                    sf_symbols[sf,nsymbol,fidx] = self.pilots[port_id // 2][sf % 10][pidx]
                    fidx += c.LTE_NRE // 2
        return sf_symbols
    
    def toIdx(self,c):
        a = c * (1+1j)*np.sqrt(0.5)
        if round(a.real) == 1:
            return 0
        elif round(a.real) == -1:
            return 1
        elif round(a.imag) == 1:
            return 2
        else:
            return 3

    def remap(self,idx):
        if idx>=36:
            return idx+1-36
        else:
            return 128+idx-36

    def cs_table(self,port_id):
        r = []
        for i in range(4):
            r.append([])
        for sf in range(c.LTE_NOF_SF_X_FRAME):        
            for l in range(c.refsignal_cs_nof_symbols(port_id)):
                nsymbol = c.refsignal_cs_nsymbol(l, self.cp, port_id)
                fidx = (c.refsignal_cs_v(port_id, l) + (self.id % 6)) % 6
                for i in range(2 * self.nof_prb):
                    pidx = c.LTE_REFSIGNAL_PILOT_IDX(i, l, self.nof_prb)
                    idx = self.toIdx(self.pilots[port_id // 2][sf % 10][pidx])
                    r[idx].append((sf,nsymbol,self.remap(fidx)))
                    fidx += c.LTE_NRE // 2
        return r

    def cs_tx(self,port_id):
        sf_symbols = self.cs_put(port_id)
        # print("ofdm_tx", sf_symbols.shape,self.nof_prb)
        return Phy.ofdm_tx(sf_symbols,self.nof_prb) 

class PSS:
    zc_map = [25,29,34]
    def __init__(self,t):
        ofd = np.exp(-1j*np.pi*self.zc_map[t]/63*np.arange(63)*np.arange(1,64))
        self.fd = np.zeros((128,),dtype=complex)
        self.fd[1:32] = ofd[32:]
        self.fd[97:]  = ofd[:31]
        self.ofd = np.zeros((62,),dtype=complex)
        self.ofd[:31] = ofd[:31]
        self.ofd[31:] = ofd[32:]

        td = np.fft.ifft(self.fd) 
        self.td = np.zeros((137,),dtype=complex)
        self.td[9:] = np.conj(td[:])
        self.td[:9] = np.conj(td[-9:])
        self.tlen = 137

    def xcorr_ft(self,rfd,sset,fc=874.2e6,sr=1.92e6):
        dlen = rfd.shape[0]-100
        fn = dlen//19200
        halffl = 9600
        r = np.zeros((len(sset),halffl))
        for foi in range(len(sset)):
            f = sset[foi]
            drc = 1+f/fc
            dfps = self.td*np.exp(-1j*np.pi*2.*f/sr*np.arange(len(self.td)))
            for i in range(halffl):
                for j in range(2*fn):
                    ppos = int((i+(j*halffl))/drc)
                    x = np.dot(rfd[ppos:ppos+self.tlen],dfps)
                    r[foi,i] += np.abs(x)
        return r

class LazySSS:
    def __init__(self):
        self.sss = [None] * 1008
    def get(self,id_1,id_2,slot):
        inx = id_1*6+id_2*2+slot
        if self.sss[inx]==None:
            self.sss[inx] = SSS(id_1,id_2,slot)
        return self.sss[inx]
    
class SSS:
    sss = LazySSS()
    def __init__(self,id_1,id_2,slot):
        qp = id_1//30
        q  = (id_1+qp*(qp+1)//2)//30
        mp = id_1+q*(q+1)//2
        m0 = mp%31
        m1 = (m0+mp//31+1)%31
        s_td = 1-2*np.array([0,0,0,0,1,0,0,1,0,1,1,0,0,1,1,1,1,1,0,0,0,1,1,0,1,1,1,0,1,0,1])
        c_td = 1-2*np.array([0,0,0,0,1,0,1,0,1,1,1,0,1,1,0,0,0,1,1,1,1,1,0,0,1,1,0,1,0,0,1])
        z_td = 1-2*np.array([0,0,0,0,1,1,1,0,0,1,1,0,1,1,1,1,1,0,1,0,0,0,1,0,0,1,0,1,0,1,1])
        s0_m0 = np.array([ s_td[(i+m0)%31] for i in range(31)])
        s1_m1 = np.array([ s_td[(i+m1)%31] for i in range(31)])
        c0    = np.array([ c_td[(i+id_2)%31] for i in range(31)])
        c1    = np.array([ c_td[(i+id_2+3)%31] for i in range(31)])
        z1_m0 = np.array([ z_td[(i+m0%8)%31] for i in range(31)])
        z1_m1 = np.array([ z_td[(i+m1%8)%31] for i in range(31)])
        if slot==0:
            ssc2 = s1_m1*c1*z1_m0
            ssc1 = s0_m0*c0
        else:
            ssc2 = s0_m0*c1*z1_m1
            ssc1 = s1_m1*c0
        self.ofd = np.zeros(62)
        self.ofd[::2] = ssc1[:]
        self.ofd[1::2] = ssc2[:]
        self.fd = np.zeros((128,),dtype=complex)
        self.fd[1:32] = self.ofd[31:]
        self.fd[97:]  = self.ofd[:31]  
        
class refsignal(Cell):
    offset = [10,558,970,1518]
    symbol_idx = [0,4,7,11]
    def __init__(self,id,nof_prb=6,cp="normal"):
        Cell.__init__(self,id,nof_prb,cp)
        self.symbol_sz = c.LTE_SYMBOL_SZ[self.nof_prb]
        self.sf_len = 1920
        
        self.ref = np.zeros((self.nof_ports,c.LTE_NOF_SF_X_FRAME,self.sf_len),dtype=complex)
        self.frame_ref_freq = np.zeros((self.nof_ports,c.LTE_NOF_SF_X_FRAME*self.sf_len),dtype=complex)
        
        for p in range(self.nof_ports):
            self.ref[p,:,:] = self.cs_tx(p)
        for p in range(self.nof_ports):
            self.frame_ref_freq[p,:] = np.fft.fft(self.frame_refsignal(p))

    def xcorr(self,port_id,buf,slot,peak):
        ch = np.zeros(128,dtype=complex)
        for off in self.offset:
            o = off+peak
            r = self.ref[port_id][slot,off:off+self.symbol_sz]
            for i in range(-10,64):
                d = buf[o+i:o+i+self.symbol_sz]
                ch[i+64] += np.dot(d,np.conj(r))
        return ch

    def xcorr_slot(self,port_id,buf,slot,peak):
        ch = np.zeros(128,dtype=complex)
        r = self.ref[port_id][slot,:]
        for i in range(-10,64):
            d = buf[peak+i:peak+i+c.LTE_SF_LEN(self.symbol_sz)]
            ch[i+64] += np.dot(d,np.conj(r))
        return ch

    def pos0(self,port_id,buf,peak):
        res = np.zeros((c.LTE_NOF_SF_X_FRAME,),dtype=complex)
        sf_idx_init = (2 * c.LTE_NOF_SF_X_FRAME - peak // self.sf_len) % c.LTE_NOF_SF_X_FRAME
        n      = peak % self.sf_len
        for _slot in range(c.LTE_NOF_SF_X_FRAME):
            slot = (sf_idx_init+_slot)%c.LTE_NOF_SF_X_FRAME
            for off in self.offset:
                d = buf[off+n:off+n+self.symbol_sz]
                r = self.ref[port_id][slot,off:off+self.symbol_sz]
                res[_slot] += np.dot(d,np.conj(r))
            n += self.sf_len
        return res

    def frame(self,port_id,buf,peak):
        ch = np.zeros(128,dtype=complex)
        sf_idx_init = (2 * 10 - peak // self.sf_len) % 10
        n      = peak % self.sf_len
        for _slot in range(10):
            slot = (sf_idx_init+_slot)%10
            ch += self.xcorr_slot(port_id,buf,slot,n)
            n += self.sf_len
        return ch

    def symbol_of_frame(self,port_id,buf,peak):
        ch = np.zeros(19200,dtype=complex)
        sf_idx_init = (2 * 10 - peak // self.sf_len) % 10
        n      = peak % self.sf_len
        for _slot in range(10):
            slot = (sf_idx_init+_slot)%10
            sym =self.symbol_of_slot(port_id,buf,slot,n)
            for i in range(4):
                ch[_slot*1920+self.offset[i]] = sym[i] 
            n += self.sf_len
        return ch
    
    def symbol_of_slot(self,port_id,buf,slot,peak):
        ch = np.zeros(4,dtype=complex)
        r = self.ref[port_id][slot,:]
        for i in range(4):
            off = self.offset[i]
            d = buf[off+peak:off+peak+self.symbol_sz]
            r = self.ref[port_id][slot,off:off+self.symbol_sz]
            ch[i] = np.dot(d,np.conj(r))
        return ch

    def afc(self,buf,port,peak):
        l = len(buf)
        p = (l-peak)//19200
        r = np.zeros(p*10-10,dtype=complex)
        for frame in range(p-1):
            r[frame*10:frame*10+10] = self.pos0(port,buf[frame*19200:frame*19200+19200*2],peak)
        return r

    def frame_refsignal(self,port_id):
        r = np.zeros(19200,dtype=complex)
        for slot in range(10):
            o = slot*1920 
            r[o:o+1920] = self.ref[port_id][slot,:]
        return r

    def find_freq(self,buf,port_id,frange):
        d = buf[:19200]
        fd = np.fft.fft(d)
        flen = len(frange)
        txr = np.zeros((flen,19200))
        for i in range(flen):
            s = frange[i]
            fr = np.zeros(19200,dtype=complex)
            if s !=0:
                fr[s:] = self.frame_ref_freq[port_id][:-s]
                fr[:s] = self.frame_ref_freq[port_id][-s:]
            else:
                fr[:] = self.frame_ref_freq[port_id]
            xr = fd*np.conj(fr)
            txr[i,:] = np.abs(np.fft.ifft(xr))
        peak = txr.argmax()
        df = peak//19200
        peak = peak%19200
        return txr[df,peak],port_id,peak,frange[df]

    def find_peak_slot(self,buf,port_id):
        d = buf[:19200]
        fd = np.fft.fft(d)
        reft = self.frame_refsignal(port_id)
        txr = np.zeros((19200,))
        for slot in range(10):
            slott = np.zeros(reft.shape,dtype=np.complex64)
            slott[slot*1920:slot*1920+1920] = reft[slot*1920:slot*1920+1920]
            reff = np.fft.fft(slott)
            xr = fd*np.conj(reff)
            txr += np.abs(np.fft.ifft(xr))
        peak = txr.argmax()
        return txr[peak],peak,port_id
    
    def find_peak(self,buf,port_id):
        d = buf[:19200]
        fd = np.fft.fft(d)
        xr = fd*np.conj(self.frame_ref_freq[port_id])
        txr = np.abs(np.fft.ifft(xr))
        peak = txr.argmax()
        return txr[peak],peak,port_id

    def find_all_peak(self,buf):
        r = [self.find_peak(buf,port_id) for port_id in range(self.nof_ports)]
        r.sort(reverse=True)
        return r

    def find_all_freq(self,buf,frange):
        r = [self.find_freq(buf,port_id,frange) for port_id in range(self.nof_ports)]
        r.sort(reverse=True)
        return r

    @staticmethod
    def all_peak(buf):
        r = []
        for cid in range(504):
            for port in range(4):
                ref = refsignal(cid)
                a,peak,_ = ref.find_peak(buf,port)
                r.append((a,cid,port,peak))
        r.sort(reverse=True)
        return r
    @staticmethod
    def all_peak_slot(buf):
        r = []
        for cid in range(504):
            for port in range(4):
                ref = refsignal(cid)
                a,peak,_ = ref.find_peak_slot(buf,port)
                r.append((a,cid,port,peak))
        r.sort(reverse=True)
        return r
class LazyReferences:
    def __init__(self):
        self.refs = [None] * 504
    
    def __getitem__(self, idx):
        if self.refs[idx]==None:
            self.refs[idx] = refsignal(idx)
        return self.refs[idx]
    

